#*-* utf-8 -*-

import numpy as np;
import scipy.sparse as sp
from scipy import linalg
import time
import threading
import os
from multiprocessing import Process, Manager
import utils
import sys
dx = 112896
dy = 112896
d = 100

beta = 1
gamma = 1

thetax = 0.6
thetay  = 0.5

converge_threshhold = 1;

#For this time, Return csc_matrix
def getColumnVector(s, d):
    row = [];
    column = [];
    data = [];
    li = s.strip().split(' ');
    for l in li:
        pair = l.strip().split(':');
        index = int(pair[0]);
        value = float(pair[1]);
        row.append(index-1);
        column.append(0);
        data.append(value)
    return sp.csc_matrix((data,(row, col)), shape=(d,1));

def saveModel(Lx, Ly):
    print 'Saving Model.....';
   
    utils.save_sparse_csc('Lxosx.txt.npz', Lx);
    utils.save_sparse_csc('Lyosx.txt.npz', Ly);
    

def calcC(row, theta):
    t = np.linalg.norm(row);
    if t == 0:
        return 0;
    return theta/t;

def checkConverge(LPx,LPy, Lx, Ly):
    t = np.linalg.norm((Lx-LPx).toarray()) + np.linalg.norm((Ly-LPy).toarray());
    if t < converge_threshhold:
        return (True, t);
    else:
        return (False, t);

def compute(x, t):
    return np.sign(x) * max(np.abs(x)-t, 0);

def updateLxu(slx, Lx, pid, start, end ,debug=False):
    #print 'pid: %d, start: %d, end: %d' % (pid, start, end);
    compute_lxu = np.frompyfunc(lambda x: compute(x, beta), 1, 1);
    for u in range(start, end):
        if debug:
            start = time.time();
        wu = WX.getcol(u);

        Lx[u] = compute_lxu(wu.transpose().todense());

        cu = calcC(Lx[u].todense(), thetax);

        Lx[u] = cu * Lx[u];

        if debug:
            print  u, time.time() -start;
    slx.append((Lx, start, end));

def updateLyv(sly, Ly, pid, start, end, debug=False):
    #print 'pid: %d, start: %d, end: %d' % (pid, start, end);
    compute_lyv = np.frompyfunc(lambda x:compute(x, gamma), 1, 1);
    for v in range(start, end):
        if debug:
            start = time.time();
        nv = NY.getcol(v);
        Ly[v] = compute_lyv(nv.transpose().todense());
        cv = calcC(Ly[v].todense(), thetay);
        Ly[v] = cv * Ly[v];

        if debug:
            print v,time.time() -start;
    sly.append((Ly, start, end));

def update(M, sl):
    t = 0;
    for triple in sl:
        R, s, e = triple;
        #t += linalg.norm((R[s:e].tocsr() - M[s:e].tocsr()).todense())
        M[s:e] = R[s:e];
        
    return t;
def RMLS():
    global beta;
    global gamma;
    t = -1;
    count = 0;
    global WX, NY, Lx, Ly;
    K = 20;
    Lx = sp.rand(dx, d, density=0.01, format='csr', dtype=np.float64);
    Ly = sp.rand(dy, d, density=0.01, format='csr', dtype=np.float64);
    manager = Manager();
    T = 0;
    while True:
        T += 1;
        LPx = Lx.copy();
        LPy = Ly.copy();
        count += 1;
        print 'This is %d iteration, last converge is %f' %(count, t);
        start = time.time();
        WX = Ly.transpose().dot(W);
        #print 'WX', WX;
        ranx = dx/K;

        processListx = [];
        u = 0;
        slx = manager.list();
        for pid in range(K-1):
            processListx.append(Process(target=updateLxu, args=(slx, Lx, pid, u, u+ranx)));
            u = u + ranx;
        processListx.append(Process(target=updateLxu, args=(slx, Lx, K-1, u, dx, False)));

        for pid in range(K):
            processListx[pid].start();
        for pid in range(K):
            processListx[pid].join();
    
        norm =  update(Lx, slx);
 
        print 'Update Lx cost: ' + str(time.time() -start); 
    
        start = time.time();
        NY = Lx.transpose().dot(N);
        #Update Ly
        rany = dy/K;
        processListy = [];
        v  = 0;
        sly  = manager.list();
        for pid in range(K-1):
            processListy.append(Process(target=updateLyv, args=(sly,Ly,pid, v, v+rany)));
            v = v + rany;
        processListy.append(Process(target=updateLyv, args=(sly, Ly, K-1, v, dy, False)));

        for pid in range(K):
            processListy[pid].start();
        for pid in range(K):
            processListy[pid].join();

        norm +=  update(Ly, sly)    
        print 'Update Ly cost: ' + str(time.time() - start);

        #print LPx, Lx;
        flag, t = checkConverge(LPx, LPy, Lx, Ly);
        #flag = False;
        #if norm < converge_threshhold:
        #    flag = True;
        print 'Converge: ', flag,t;
        if flag:
            break;
        if T >= 25:
            break;


    saveModel(Lx, Ly);

if __name__ == '__main__':
    #featureDataFile = 'train.set';
    #batchPreprocessing(featureDataFile);
    try:
        global W;
        global N;
        W = sp.csc_matrix((dy, dx), dtype = np.float64);
        N = sp.csc_matrix((dx,dy), dtype = np.float64);
        W = utils.load_sparse_csc('../Wxs.txt.npz');
        N = utils.load_sparse_csc('../Nxs.txt.npz');
        print W.shape;
        print N.shape;
        # print 'XXXXXXXXXXXXXXXXXXXXXX', type(linalg.norm(W))
        #print 'XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX'
        # print W, N
        RMLS(); 
    except KeyboardInterrupt:
        saveModel(Lx, Ly);
        sys.exit(0);
